package controller

import (
	"context"
	"errors"
	"fmt"
	"ollama-go/dto"
	"ollama-go/query"
	"ollama-go/service"
	"ollama-go/toolkit"
	"strings"

	"github.com/androidsr/sc-go/model"
	"github.com/androidsr/sc-go/sgin"
	"github.com/gin-gonic/gin"
	"github.com/tmc/langchaingo/schema"
)

func init() {
	sgin.AddRouter(OllamaEmbeddingController{
		ollamaEmbeddingService:  service.NewOllamaEmbeddingService(),
		ollamaCollectionService: service.NewOllamaCollectionService(),
	})
}

type OllamaEmbeddingController struct {
	ollamaEmbeddingService  *service.OllamaEmbeddingService
	ollamaCollectionService *service.OllamaCollectionService
}

// 查询
// @Router [get] [/ollamaembedding/:id]
func (m OllamaEmbeddingController) Get(c *gin.Context) any {
	id := c.Param("id")
	if id == "" {
		return model.NewFailDefault("ID不能为空")
	}
	return m.ollamaEmbeddingService.Get(id)
}

// 增加
// @Router [post] [/ollamaembedding]
func (m OllamaEmbeddingController) Save(c *gin.Context, dto *dto.OllamaEmbeddingDTO) any {
	dto.TenantId = c.GetString(USER)
	err := m.ollamaEmbeddingService.Save(dto)
	if err != nil {
		return model.NewFailDefault(err.Error())
	} else {
		return model.NewOK("创建成功")
	}
}

// 修改
// @Router [put] [/ollamaembedding/edit]
func (m OllamaEmbeddingController) UpdateById(c *gin.Context, dto *dto.OllamaEmbeddingDTO) any {
	if dto.Uuid == "" {
		return model.NewFailDefault("ID不能为空")
	}
	dto.TenantId = c.GetString(USER)
	return m.ollamaEmbeddingService.UpdateById(dto)

}

// 删除
// @Router [delete] [/ollamaembedding/:id]
func (m OllamaEmbeddingController) DeleteById(c *gin.Context) any {
	id := c.Param("id")
	if id == "" {
		return model.NewFailDefault("ID不能为空")
	}
	return m.ollamaEmbeddingService.DeleteById(id)
}

// 查询列表
// @Router [post] [/ollamaembedding/page]
func (m OllamaEmbeddingController) QueryPage(c *gin.Context, query *query.OllamaEmbeddingQueryDTO) any {
	query.TenantId = c.GetString(USER)
	return m.ollamaEmbeddingService.QueryPage(query)
}

// 分页下拉
// @Router [post] [/ollamaembedding/list]
func (m OllamaEmbeddingController) QuerySelect(c *gin.Context, query *model.SelectQueryDTO) any {
	return m.ollamaEmbeddingService.QueryList(query)
}

// @Router [post] [/ollamaembedding/upload/:collectionName]
func (m OllamaEmbeddingController) Upload(c *gin.Context) any {
	file, _ := c.FormFile("file")
	collectionName := c.Param("collectionName")
	tenantId := m.ollamaCollectionService.GetTenantId(collectionName)
	cmetadata := make(map[string]any, 0)
	cmetadata["tenantId"] = tenantId
	return toolkit.PG.ReadFormFile(c, cmetadata, file)
}

// @Router [post] [/ollamaembedding/uploads/:collectionName]
func (m OllamaEmbeddingController) Uploads(c *gin.Context) any {
	form, _ := c.MultipartForm()
	files := form.File["files"]
	errors := make([]string, 0)
	collectionName := c.Param("collectionName")
	tenantId := m.ollamaCollectionService.GetTenantId(collectionName)
	cmetadata := make(map[string]any, 0)
	cmetadata["tenantId"] = tenantId
	fmt.Println(cmetadata)
	for _, file := range files {
		result := toolkit.PG.ReadFormFile(c, cmetadata, file)
		if result.Code != 200 {
			errors = append(errors, file.Filename)
		}
	}
	if len(errors) != 0 {
		return model.NewFail(55110, strings.Join(errors, ","))
	}
	return model.NewOK("处理成功")
}

// @Router [post] [/ollamaembedding/deleteFilename]
func (m OllamaEmbeddingController) DeleteFilename(c *gin.Context) any {
	filename := c.Query("filename")
	collectionId := c.Query("collectionId")
	if filename == "" {
		return model.NewFailDefault("删除文件不能为空")
	}
	err := m.ollamaEmbeddingService.DeleteByFilename(collectionId, filename)
	if err != nil {
		return model.NewFailDefault(err.Error())
	}
	return model.NewOK("删除成功")
}

// 增加
// @Router [post] [/ollamaembedding/addDocument]
func (m OllamaEmbeddingController) AddDocument(c *gin.Context, dto *dto.AddDocumentDTO) any {
	if dto.CollectionName == "" {
		collection := service.NewOllamaCollectionService().Get(dto.CollectionId)
		if collection == nil {
			return errors.New("集合不存在")
		}
		dto.CollectionName = collection.Name
	}
	err := toolkit.PG.SetStore(context.Background(), dto.CollectionName).AddDocument(context.Background(), []schema.Document{{
		PageContent: dto.Document,
		Metadata:    dto.Metadata,
	}})
	if err != nil {
		return model.NewFailDefault(err.Error())
	}
	return model.NewOK("保存成功")
}

// @Router [post] [/ollamaembedding/deleteMetaData]
func (m OllamaEmbeddingController) DeleteMetaData(c *gin.Context) any {
	params := make(map[string]interface{}, 0)
	err := c.BindJSON(&params)
	if err != nil {
		return model.NewFailDefault(err.Error())
	}
	err = m.ollamaEmbeddingService.DeleteMetaData(params)
	if err != nil {
		return model.NewFailDefault(err.Error())
	}
	return model.NewOK("删除成功")
}

// @Router [post] [/ollamaembedding/search]
func (m OllamaEmbeddingController) Search(c *gin.Context) any {
	input := &toolkit.SearchDTO{}
	err := c.BindJSON(input)
	if err != nil {
		return model.NewFailDefault(err.Error())
	}
	docs, err := toolkit.PG.SetStore(c.Request.Context(), input.CollectionName).Search(c.Request.Context(), input)
	if err != nil {
		sendMessage(c.Writer, []byte(err.Error()))
		return model.NewFailDefault(err.Error())
	}
	result := make([]map[string]interface{}, 0)
	var item map[string]interface{}
	for _, v := range docs {
		item = make(map[string]interface{}, 0)
		item["document"] = v.PageContent
		item["score"] = 1 - v.Score
		item["metadata"] = v.Metadata
		result = append(result, item)
	}
	return model.NewOK(result)
}
